import pandas as pd


def cal_emaN(days, i, close, lastemaN):
    if i == 1:
        return close[i - 1]
    return lastemaN[i - 2] * (days - 1) / (days + 1) + close[i - 1] * 2 / (days + 1)


def cal_deaN(ma, dea, diff, i):
    if i == 1:
        return 0
    return dea[i - 2] * (ma - 1) / (ma + 1) + diff[i - 1] * 2 / (ma + 1)


def calMACD(df, fast_ma, slow_ma, ma):
    sort_df = df.sort_index(ascending=True)
    ema_fast, ema_slow, DIFF, DEA9, MACDN = [], [], [], [], []
    close = sort_df['close']

    for i in range(1, close.count() + 1):
        ema12 = cal_emaN(fast_ma, i, close, ema_fast)
        ema_fast.append(ema12)
        ema26 = cal_emaN(slow_ma, i, close, ema_slow)
        ema_slow.append(ema26)
        diff = ema12 - ema26
        DIFF.append(diff)
        dea = cal_deaN(ma, DEA9, DIFF, i)
        DEA9.append(dea)
        MACDN.append(2 * (diff - dea))

    sort_df[f'ema{fast_ma}'] = ema_fast
    sort_df[f'ema{slow_ma}'] = ema_slow
    sort_df['diff'] = DIFF
    sort_df[f'dea{ma}'] = DEA9
    sort_df['macd'] = MACDN
    return sort_df


if __name__ == '__main__':
    df = pd.read_csv('data_qfq/sz300347.csv')
    df_ = calMACD(df, 12, 26, 9)
    print(df_)
    